# -*- coding: UTF-8 -*-

# %% [markdown]
#
import numpy as np
import matplotlib.pyplot as plt
import control as ct
import sympy as sp

def ztrans2diffexpr(tf_z: sp.core.mul.Mul):
    degree_num = sp.degree(sp.numer(tf_z), z)
    degree_den = sp.degree(sp.denom(tf_z), z)
    tf_z_num = sp.numer(tf_z)  # /z**degree_num
    tf_z_num = sp.simplify(tf_z_num)
    tf_z_num_coeff = [tf_z_num.coeff(z, i)
                                for i in range(degree_num+1)]
    # tf_z_num_coeff.reverse()
    tf_z_den = sp.denom(tf_z)  # /z**degree_den
    tf_z_den = sp.simplify(tf_z_den)
    tf_z_den_coeff = [tf_z_den.coeff(z, i)
                                for i in range(degree_den+1)]
    # tf_z_den_coeff.reverse()

    k = tf_z_den_coeff[-1]
    # m_ = [x/k for x in tf_z_num_coeff]
    # n_ = [-x/k for x in tf_z_den_coeff[:-1]]
    # m_.reverse()
    # n_.reverse()
    m = (np.array(tf_z_num_coeff)/k)[::-1]
    n = (-np.array(tf_z_den_coeff)[:-1]/k)[::-1]

    return m, n


def bi_filter_generate_gain(Ts_, fc_, d_depth_, zeta_0N_, zeta_0D_):

    Ts, fc, d_depth = sp.symbols('Ts f_c d_depth')

    s = sp.symbols('s')
    z = sp.symbols('z')
    w_0N, zeta_0N, w_0D, zeta_0D = sp.symbols('w_0N zeta_N w_0D zeta_D')
    f_0N, f_0D = sp.symbols('f_0N f_0D')

    equ = [sp.Eq(d_depth, (f_0D/f_0N)**2),
           sp.Eq(f_0D, fc),
           sp.Eq(w_0N, 2*sp.pi*f_0N),
           sp.Eq(w_0D, 2*sp.pi*f_0D)]
    sol = sp.solve(equ, w_0N, w_0D, f_0N, f_0D)
    w_0N = sol[0][0]
    w_0D = sol[0][1]
    f_0N = sol[0][2]
    f_0D = sol[0][3]

    sys = d_depth*(s**2 + 2*zeta_0N*w_0N*s + w_0N**2) / \
        (s**2 + 2*zeta_0D*w_0D*s + w_0D**2)

    sys = sys.subs(s, 2/Ts*(z-1)/(z+1))
    sys_tmp = sys

    sys_tmp = sp.simplify(sys_tmp)
    sys_tmp = sp.expand(sys_tmp, frac=True)
    sys_tmp = sp.collect(sys_tmp, z)
    sys_tmp_value = sys_tmp
    sys_tmp_value = sys_tmp.subs(
        {Ts: Ts_, d_depth: d_depth_, fc: fc_, zeta_0N: zeta_0N_, zeta_0D: zeta_0D_})
    # sys_tmp_value = sys_tmp.subs({Ts:1e-4, d_depth:0.5, f_width:3, f0:15})

    m, n = ztrans2diffexpr(sys_tmp_value)
    return m, n
    
def bi_filter(x, x_last1, x_last2, y_last1, y_last2, m, n):

    y = (m[0]*x + m[1]*x_last1 + m[2]*x_last2) + \
            (n[0]*y_last1 + n[1]*y_last2)

    return y

if __name__ == "__main__":
    pass
    # %% [markdown]
    #
    x, y = sp.symbols('x y')
    z = sp.symbols('z')

    # %% [markdown]
    # form
    Ts, fc, d_depth, zeta_0N, zeta_0D = sp.symbols('Ts, f_c, d_depth, zeta_0N, zeta_0D')
    m, n = bi_filter_generate_gain(Ts, fc, d_depth, zeta_0N, zeta_0D)
    print('m[0]:{}'.format(m[0]))
    print('m[1]:{}'.format(m[1]))
    print('m[2]:{}'.format(m[2]))
    print('n[0]:{}'.format(n[0]))
    print('n[1]:{}'.format(n[1]))
    
    
    # %% [markdown]
    # data test
    Ts0, d_deepth0, f_width0, f00 = sp.symbols('Ts0 d_deepth0 f_width0 f00')
    # m, n = filter_generate_gain(Ts0, d_deepth0, f_width0, f00)
    m, n = bi_filter_generate_gain(1e-4, 50, 1/3, 0.2, 0.5)
    m = [float(m[i]) for i in range(len(m))]
    n = [float(n[i]) for i in range(len(n))]
    print(m)
    print(n)

    x, x_last1, x_last2, y_last1, y_last2 = 0, 0, 0, 0, 0

    f = 600
    t = np.arange(20000)*1e-4
    sin_data = np.sin(2*np.pi*f*t)
    # plt.scatter(t, sin_data)
    # plt.show()

    x_list = []
    y_list = []
    y = 0
    for i in range(len(t)):
        x_last2 = x_last1
        x_last1 = x
        x = sin_data[i]
        y_last2 = y_last1
        y_last1 = y
        x_list.append(x)
        y = bi_filter(x, x_last1, x_last2, y_last1, y_last2, m, n)
        y_list.append(y)

    plt.scatter(t, y_list)
    plt.show()
    print(max(x_list[10000:]))
    print(max(y_list[10000:]))